import numpy as np
import matplotlib.pyplot as plt
from tqdm import tqdm

# Function to generate a dataset based on a sine wave
def generate_dataset(size=200,timesteps=25):
    x,y = [],[]
    sin_wave = np.sin(np.arange(size))
    for step in range(sin_wave.shape[0] - timesteps):
        x.append(sin_wave[step:step + timesteps])
        y.append(sin_wave[step + timesteps])
    return np.array(x).reshape(len(y),timesteps,1), np.array(y).reshape(len(y),1)

x,y = generate_dataset()
print(x.shape,y.shape)

print(x)
print(y)
plt.plot(x[-1])
plt.plot(y)

# simple rnn class
class SimpleRNN:
    def __init__(self,x,y,hidden_units):
        self.x = x
        self.y = y
        self.hidden_units = hidden_units
        # initialize weights randomly
        self.Wx = np.random.randn(self.hidden_units,self.x.shape[2])
        self.Wh = np.random.randn(self.hidden_units,self.hidden_units)
        self.Wy = np.random.randn(self.y.shape[1],self.hidden_units)
        print(self.Wx.shape,self.Wh.shape,self.Wy.shape)

    # function for single rnn cell
    def cell(self,xt,ht_1):
        # print(xt.shape,ht_1.shape)
        # compute the hidden state
        ht = np.tanh(np.dot(self.Wx,xt.reshape(1,1)) + np.dot(self.Wh,ht_1))
        # print(ht.shape)
        # compute the output
        yt = np.dot(self.Wy,ht)
        # print(yt.shape)
        return ht,yt

    # forward pass through the network
    def forward(self,sample):
        sample_x, sample_y = self.x[sample], self.y[sample]
        # print(sample_x.shape,sample_y.shape)

        ht = np.zeros((self.hidden_units,1))
        self.hidden_states = [ht]
        self.inputs = []
        for step in range(len(sample_x)):
            ht,yt = self.cell(sample_x[step],ht)
            self.inputs.append(sample_x[step].reshape(1,1))
            self.hidden_states.append(ht)
        # print(len(self.hidden_states),self.hidden_states[-1].shape)
        self.error = yt - sample_y
        self.loss = 0.5 * self.error**2
        self.yt = yt

    # backward pass through the network(backpropagetion)
    def backward(self):
        n = len(self.inputs)
        # compute the gradient of output
        dyt = self.error
        # compute the gradient of output weights
        dWy = np.dot(dyt,self.hidden_states[-1].T)

        # compute the gradient of the hidden state
        dht = np.dot(dyt,self.Wy).T
        dWx = np.zeros(self.Wx.shape)
        dWh = np.zeros(self.Wh.shape)

        # backpropage through time
        for step in reversed(range(n)):
            temp = (1 - self.hidden_states[step+1]**2) * dht
            dWx += np.dot(temp,self.inputs[step].T)
            dWh += np.dot(temp,self.hidden_states[step].T)
            dht = np.dot(self.Wh,temp)

        # clip to prevent exloding gradients
        dWy = np.clip(dWy,-1,1)
        dWx = np.clip(dWx,-1,1)
        dWh = np.clip(dWh,-1,1)

        # update weights
        self.Wy -= self.lr * dWy
        self.Wx -= self.lr * dWx
        self.Wh -= self.lr * dWh

    # training function
    def train(self,epochs,learning_rate):
        self.Ovr_loss = []
        self.lr = learning_rate
        for epoch in tqdm(range(epochs)):

            for sample in range(self.x.shape[0]):
                self.forward(sample)
                self.backward()
            self.Ovr_loss.append(np.squeeze(self.loss / self.x.shape[0]))
            print(f"epoch: {epoch}  ovr_loss: f{self.Ovr_loss}")
            self.loss = 0

    # testing function
    def test(self,x,y):
        self.x = x
        self.y = y
        self.outputs = []
        for sample in range(len(x)):
            self.forward(sample)
            self.outputs.append(self.yt)

# generate dataset
x,y = generate_dataset()
x_test,y_test = generate_dataset(300)
x_test = x_test[250:]
y_test = y_test[250:]

# create simpleRNN instance and train
rnn = SimpleRNN(x,y,100)
rnn.train(25,1e-2)

# test the trained simplernn
rnn.test(x_test,y_test)

# plot the results
plt.tight_layout()
plt.figure(dpi=120)
plt.subplot(121)
print('----------loss曲线-------------')
plt.plot(rnn.Ovr_loss)
plt.subplot(122)
print('----------预测值和真实值的对比-------------')
plt.plot([i for i in range(len(x_test))],y_test,np.array(rnn.outputs).reshape(y_test.shape))
print("---------------------")
print(x.shape)
print(y.shape)








